# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon
import pydevd
from collections import namedtuple
import concurrent.futures
import csv
import datetime
import os
import pathlib
import sys
import traceback

from absl import (flags, app)

from coin.base.book.book_builder_base import ThinBookError
from coin.base.param_util import to_list
from coin.exchange.binance.kr_rest.product import BinanceProduct
from coin.exchange.bitflyer_v1.kr_rest.futures_product import (BitflyerFuturesProduct)
from coin.exchange.bitfinex_v2.kr_rest.product import BitfinexProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.gdax.kr_rest.product import GdaxProduct
from coin.exchange.huobi.kr_rest.product import HuobiProduct
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.bithumb.kr_rest.product import BithumbProduct
from coin.exchange.okex_futures.kr_rest.futures_product import (OkexFuturesProduct)
from coin.exchange.kraken.kr_rest.product import KrakenProduct
from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.dumper_base import run_from_archive
from coin.support.feed_tool.feed_stats.logic.util import *
from coin.exchange.util.feed_checker import get_default_feed_checker_config

FLAGS = flags.FLAGS

CsvWriter = namedtuple('CsvWriter', ['writer', 'file', 'notional'])


def get_notional_qty(product):
  if isinstance(product, BitmexFuturesProduct):
    return 10000
  elif isinstance(product, OkexFuturesProduct):
    return 100
  elif isinstance(product,
                  (BitflyerFuturesProduct,
                   BinanceProduct,
                   BithumbProduct,
                   BitfinexProduct,
                   GdaxProduct,
                   HuobiProduct,
                   OkexProduct,
                   KrakenProduct)):
    return 1.2
  else:
    raise ValueError()


def get_notional_amount(product):
  return 1000
  """
  notional_amount_dict = {
      BitfinexProduct.FromStr('BCH-USDT'): 1000,
      BitfinexProduct.FromStr('BTC-USDT'): 1000,
      BitfinexProduct.FromStr('EOS-USDT'): 1000,
      BitfinexProduct.FromStr('ETH-USDT'): 1000,
      BitfinexProduct.FromStr('IOTA-USDT'): 1000,
      BitfinexProduct.FromStr('LTC-USDT'): 1000,
      BitfinexProduct.FromStr('XRP-USDT'): 1000,
      HuobiProduct.FromStr('BCH-USDT'): 1000,
      HuobiProduct.FromStr('BTC-USDT'): 1000,
      HuobiProduct.FromStr('ETH-USDT'): 1000,
      HuobiProduct.FromStr('EOS-USDT'): 1000,
      HuobiProduct.FromStr('IOTA-USDT'): 1000,
      HuobiProduct.FromStr('LTC-USDT'): 1000,
      HuobiProduct.FromStr('XRP-USDT'): 1000,
      OkexProduct.FromStr('BCH-USDT'): 1000,
      OkexProduct.FromStr('BTC-USDT'): 1000,
      OkexProduct.FromStr('EOS-USDT'): 1000,
      OkexProduct.FromStr('ETH-USDT'): 1000,
      OkexProduct.FromStr('IOTA-USDT'): 1000,
      OkexProduct.FromStr('LTC-USDT'): 1000,
      OkexProduct.FromStr('XRP-USDT'): 1000,

  }
  return notional_amount_dict[product]
  """


class DumpStrategy(object):
  def __init__(self, machine, trading_date, products, csv_root, book_levels, notional_by_qty=True):
    self._machine = machine
    self._trading_date = trading_date
    assert isinstance(products, (list, tuple))
    self._products = products
    self._csv_root = csv_root
    self._notional_by_qty = notional_by_qty
    self._book_levels = book_levels
    pathlib.Path(csv_root).mkdir(parents=True, exist_ok=True)
    self._csv_writers = {}

  def prepare_csv_file(self, product):
    path_name = os.path.join(self._csv_root,
                             self._machine,
                             self._trading_date.strftime('%Y%m%d'),
                             product.exchange)
    if not os.path.isdir(path_name):
      os.makedirs(path_name, exist_ok=True)

    file_name = '%s/%s.csv' % (path_name, product.symbol)
    csv_file = open(file_name, 'w', encoding='utf-8')
    writer = csv.writer(csv_file, dialect='excel')
    csv_header = ['timestamp', 'true_bid', 'true_ask']
    for level in range(self._book_levels):
      csv_header += ['bid%d' % level, 'bid_qty%d' % level, 'ask%d' % level, 'ask_qty%d' % level]
    writer.writerow(csv_header)
    csv_file.flush()
    return writer, csv_file

  def _get_notional_value(self, product):
    if self._notional_by_qty:
      notional = get_notional_qty(product)
    else:
      notional = get_notional_amount(product)
    return notional

  def _get_true_book_func(self, book):
    bid_func = book.get_bid_price_by_qty if self._notional_by_qty else \
               book.get_bid_price_by_amount
    ask_func = book.get_ask_price_by_qty if self._notional_by_qty else \
               book.get_ask_price_by_amount
    return (bid_func, ask_func)

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      if book_builder_name == get_subscription_key(product):
        book_builder.subscribe_book(product, self.on_book)
        book_builder.subscribe_trade(product, self.on_trade)
        exchange, symbol = to_book_builder_exchange_symbol(product)
        writer, csv_file = self.prepare_csv_file(product)
        self._csv_writers[(exchange, symbol)] = CsvWriter(writer,
                                                          csv_file,
                                                          self._get_notional_value(product))

  def on_book(self, book):
    writer = self._csv_writers.get((book.exchange, book.symbol), None)
    assert writer is not None

    bids = [[None, None]] * self._book_levels
    level = 0
    for key, value in book.iteritems_bid():
      assert key == value.price
      assert value.qty > 0
      bids[level] = [value.price, value.qty]
      level += 1
      if level >= self._book_levels:
        break
    asks = [[None, None]] * self._book_levels
    level = 0
    for key, value in book.iteritems_ask():
      assert key == value.price
      assert value.qty > 0
      asks[level] = [value.price, value.qty]
      level += 1
      if level >= self._book_levels:
        break

    try:
      bid_func, ask_func = self._get_true_book_func(book)
      csv_line = [book.timestamp, bid_func(writer.notional), ask_func(writer.notional)]
    except (IndexError, ThinBookError):
      return
    for level in range(self._book_levels):
      csv_line += [bids[level][0], bids[level][1], asks[level][0], asks[level][1]]
    writer.writer.writerow(csv_line)
    writer.file.flush()

  def on_trade(self, _):
    pass

  def close_csvs(self):
    for writer in self._csv_writers.values():
      writer.file.flush()
      writer.file.close()


def launch(products, trading_date, machine, csv_root, book_levels, hours, enable_feed_checker):
  feed_checker_config = None
  if enable_feed_checker:
    feed_checker_option = get_default_feed_checker_config(products)

  start_time = trading_date  # UTC 0 hour

  strategy = DumpStrategy(machine,
                          trading_date,
                          products,
                          csv_root,
                          book_levels,
                          notional_by_qty=False)

  run_from_archive(products,
                   strategy.on_book_reset,
                   start_time,
                   start_time + datetime.timedelta(hours=hours),
                   worker_id='2',
                   machine=[machine],
                   feed_checker_config=feed_checker_config)

  strategy.close_csvs()
  return 0


def main(argv):

  products = [
      # BitfinexProduct.FromStr('ETC-USDT'),
      # BitfinexProduct.FromStr('IOT-USDT'),
      # BitfinexProduct.FromStr('NEO-USDT'),
      # BitfinexProduct.FromStr('ETP-USDT'),
      # HuobiProduct.FromStr('ETC-USDT'),
      # HuobiProduct.FromStr('HT-USDT'),
      # HuobiProduct.FromStr('IOST-USDT'),
      # HuobiProduct.FromStr('ONT-USDT'),
      # OkexProduct.FromStr('ETC-USDT'),
      # OkexProduct.FromStr('WFEE-USDT'),
      # OkexProduct.FromStr('OKB-USDT'),
      # OkexProduct.FromStr('ONT-USDT'),
      # OkexProduct.FromStr('MITH-USDT'),
      BitfinexProduct.FromStr('BCH-USDT'),
      BitfinexProduct.FromStr('BTC-USDT'),
      BitfinexProduct.FromStr('EOS-USDT'),
      BitfinexProduct.FromStr('ETH-USDT'),
      BitfinexProduct.FromStr('LTC-USDT'),
      BitfinexProduct.FromStr('XRP-USDT'),
      HuobiProduct.FromStr('BCH-USDT'),
      HuobiProduct.FromStr('BTC-USDT'),
      HuobiProduct.FromStr('ETH-USDT'),
      HuobiProduct.FromStr('EOS-USDT'),
      HuobiProduct.FromStr('IOTA-USDT'),
      HuobiProduct.FromStr('LTC-USDT'),
      HuobiProduct.FromStr('XRP-USDT'),
      OkexProduct.FromStr('BCH-USDT'),
      OkexProduct.FromStr('BTC-USDT'),
      OkexProduct.FromStr('EOS-USDT'),
      OkexProduct.FromStr('ETH-USDT'),
      OkexProduct.FromStr('IOTA-USDT'),
      OkexProduct.FromStr('LTC-USDT'),
      OkexProduct.FromStr('XRP-USDT'),
  ]

  trading_date = FLAGS.trading_date
  machine = FLAGS.machine
  csv_root = FLAGS.csv_root
  book_levels = FLAGS.book_levels
  max_workers = FLAGS.max_workers
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert csv_root, '--csv_root must be specified.'
  print('Running for %s %s ...' % (trading_date, machine))
  sys.stdout.flush()
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')

  # pydevd.settrace('10.81.1.111', port=50000, stdoutToServer=True, stderrToServer=True)
  with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
    future_product_list = []
    for product in products:
      product = to_list(product)
      future = executor.submit(launch,
                               product,
                               trading_date,
                               machine,
                               csv_root,
                               book_levels,
                               FLAGS.hours,
                               FLAGS.enable_feed_checker)
      elem = (future, product)
      future_product_list.append(elem)
    for future, product in future_product_list:
      try:
        future.result()
      except Exception as e:
        print('Fail to dump to csv.', machine, trading_date, product, type(e), e)
        print(traceback.format_exc())


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('machine', None, 'Instance machine name.')

  flags.DEFINE_string('csv_root', None, 'Output csv files root directory.')

  flags.DEFINE_integer('book_levels', 5, 'How many levels of book to print in csv.')

  flags.DEFINE_boolean('enable_feed_checker',
                       False,
                       'Enable feed checker, default is False for performance season.')

  flags.DEFINE_integer('hours', 24, 'Time span in hours from 00:00:00 UTC.')

  flags.DEFINE_integer('max_workers', 12, 'Max number of workers.')

  app.run(main)
